﻿#include <iostream>
#include <opencv2/opencv.hpp>

// 判断颜色是否为红色
bool isRedColor(cv::Vec3b color) {
    int b = color[0];
    int g = color[1];
    int r = color[2];
    return r > 150 && g < 100 && b < 100;
}

// 判断颜色是否为蓝色
bool isBlueColor(cv::Vec3b color) {
    int b = color[0];
    int g = color[1];
    int r = color[2];
    return r < 100 && g < 100 && b > 150;
}

// 灯条识别函数
void detectLightBars(cv::Mat frame, const std::string& color1) {
    cv::Mat hsvFrame;
    cv::cvtColor(frame, hsvFrame, cv::COLOR_BGR2HSV);

    // 灰度图转化，边缘检测只能在灰度图中进行
    cv::Mat grayFrame;
    cv::cvtColor(frame, grayFrame, cv::COLOR_BGR2GRAY);

    // 自适应阈值处理
    cv::Mat binaryFrame;
    cv::adaptiveThreshold(grayFrame, binaryFrame, 255, cv::ADAPTIVE_THRESH_GAUSSIAN_C, cv::THRESH_BINARY, 11, 2);

    // 创建结构元素用于膨胀
    cv::Mat structuringElement = cv::getStructuringElement(cv::MORPH_RECT, cv::Size(5, 5));

    // 进行膨胀操作
    cv::dilate(binaryFrame, binaryFrame, structuringElement);

    // 高斯模糊
    cv::Mat blurredFrame;
    cv::GaussianBlur(grayFrame, blurredFrame, cv::Size(5, 5), 0);

    // 检测图像边缘
    cv::Mat edges;
    cv::Canny(blurredFrame, edges, 50, 150);

    // 找到边缘连接轮廓,定义轮廓列表,每个轮廓由一系列点组成
    std::vector<std::vector<cv::Point>> contours;
    cv::findContours(edges, contours, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);

    std::vector<cv::Rect> lightBars;

    for (const auto& contour : contours) {
        cv::Rect boundingRect = cv::boundingRect(contour);
        float area = cv::contourArea(contour);
        if (area > 1000) { // 调整面积阈值以过滤小噪声
            float aspectRatio = (float)boundingRect.width / boundingRect.height;
            if (aspectRatio > 0.5 && aspectRatio < 2.0) {
                for (int y = boundingRect.y; y < boundingRect.y + boundingRect.height; y++) {
                    for (int x = boundingRect.x; x < boundingRect.x + boundingRect.width; x++) {
                        if (x < frame.cols && y < frame.rows) {
                            cv::Vec3b color = frame.at<cv::Vec3b>(y, x);
                            if ((color1 == "red" && isRedColor(color)) || (color1 == "blue" && isBlueColor(color))) {
                                lightBars.push_back(boundingRect);
                                break;
                            }
                        }
                    }
                }
            }
        }
    }

    // 绘制灯条
    for (const auto& lightBar : lightBars) {
        cv::rectangle(frame, lightBar, cv::Scalar(0, 255, 0), 2);
    }

    cv::imshow("Light Bars", frame);
}

int main() {
    std::string videoPath = "C:\\Users\\aaa\\Pictures\\Screenshots\\步兵素材蓝车旋转-ev-+3.MOV";
    cv::VideoCapture video(videoPath);

    if (!video.isOpened()) {
        std::cout << "Could not open the video." << std::endl;
        return -1;
    }

    cv::Mat frame;
    std::string enermycolor = "blue";
    while (true) {
        video >> frame;
        if (frame.empty())
            break;

        detectLightBars(frame, enermycolor);

        if (cv::waitKey(30) >= 0)
            break;
    }

    video.release();
    cv::destroyAllWindows();

    return 0;
}